# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: leon

from typing import List, Dict, Tuple
from concurrent.futures import ThreadPoolExecutor
import time
import logging

from web3 import Web3, HTTPProvider, middleware

from coin.exchange.uniswap_v2.kr_rest.product import UniswapProduct
from coin.exchange.uniswap_base.constants import (
    balance_checker_deployment_address,
    PRODUCTION_RPC_URL,
    TOKEN_TO_ADDRESS,
    TOKEN_TO_DECIMAL_MULTIPLIER,
    MAX_UINT_128,
    MAX_UINT_256,
    ROUGH_ESTIMATE_GAS_LIMIT,
    BLOCKCHAIN_NAME_TO_CHAINID,
)
from coin.exchange.uniswap_v2.kr_rest.constants import (
    deployment_address,
)
from coin.exchange.uniswap_v2.kr_rest.native_private_client_utils import (
    _load_contract,
)
from coin.exchange.util.web3.signing_util import sign_and_wait


class UniswapV2NativePrivateClientV2(object):
    _blockchain_name: str
    _web3: Web3

    def __init__(
        self,
        blockchain_name='ethereum',
        web3=None,
    ):
        assert blockchain_name == 'ethereum'
        self._executor = ThreadPoolExecutor(max_workers=10)
        self._blockchain_name = blockchain_name
        if web3 is not None:
            self._web3 = web3
        else:
            self._web3 = Web3(HTTPProvider(
                PRODUCTION_RPC_URL[self._blockchain_name], request_kwargs={"timeout": 30}))
            self._web3.middleware_onion.add(middleware.time_based_cache_middleware)
            self._web3.middleware_onion.add(middleware.latest_block_based_cache_middleware)
            self._web3.middleware_onion.add(middleware.simple_cache_middleware)

    def _router(self):
        return _load_contract(
            self._web3,
            "router02",
            deployment_address["UniswapV2Router02"],
        )
    router = property(fget=_router, fset=None, fdel=None, doc=None)

    def _addr_balance_checker(self):
        return _load_contract(
            self._web3,
            "balanceChecker",
            balance_checker_deployment_address["mainnet"],
        )
    addr_balance_checker = property(
        fget=_addr_balance_checker, fset=None, fdel=None, doc=None
    )

    def get_balances(
        self,
        wallet_address: str,
        currencies: List[str],
    ) -> Dict[str, Dict[str, float]]:
        currency_addresses = [TOKEN_TO_ADDRESS['ethereum'][currency] for currency in currencies]
        balances = self.addr_balance_checker.functions.balances(
            [wallet_address], currency_addresses).call()

        currency_to_balance = {}
        for currency, balance in zip(currencies, balances):
            currency_to_balance[currency] = {
                'available': balance * TOKEN_TO_DECIMAL_MULTIPLIER.get(currency, 1e-18),
                'hold': 0,
            }
        return currency_to_balance

    def submit_order(
        self,
        wallet_address,
        symbol,
        price,
        qty,
        side,
    ) -> Tuple[bool, str or None, float or None]:
        price = float(price)
        qty = float(qty)

        p = UniswapProduct.FromStr(symbol)
        base_currency, quote_currency = p.base.currency, p.quote.currency

        trading_in_ccx = quote_currency == "CCX" or base_currency == "CCX"
        if trading_in_ccx:
            slippage = 0.2
        else:
            slippage = 0.01

        base_address = TOKEN_TO_ADDRESS['ethereum'][base_currency]
        quote_address = TOKEN_TO_ADDRESS['ethereum'][quote_currency]

        if quote_address == TOKEN_TO_ADDRESS['ethereum']["ETH"]:
            quote_address = TOKEN_TO_ADDRESS['ethereum']["WETH"]

        if base_address == TOKEN_TO_ADDRESS['ethereum']["ETH"]:
            base_address = TOKEN_TO_ADDRESS['ethereum']["WETH"]

        path = [quote_address, base_address]
        assert path[0] != path[-1]

        if (quote_address != TOKEN_TO_ADDRESS['ethereum']["WETH"]) and (base_address != TOKEN_TO_ADDRESS['ethereum']["WETH"]):
            path = [quote_address, TOKEN_TO_ADDRESS['ethereum']["WETH"], base_address]

        deadline = int(time.time() + 120)

        args, func, tx_param = None, None, {
            "nonce": None,
            "from": wallet_address,
            "type": "0x2",
            "maxFeePerGas": None,
            "maxPriorityFeePerGas": None,
            "gas": ROUGH_ESTIMATE_GAS_LIMIT,
            "chainId": BLOCKCHAIN_NAME_TO_CHAINID[self._blockchain_name],
        }
        ETH_USED = 0

        # BUY short quote long base
        # SELL short base long quote
        match side.upper():
            case 'BUY':
                # reform swap path, short quote and long base
                if path[0] != quote_address:
                    path = path[::-1]  # swap in quote and swap out base
                assert path[0] == quote_address
                assert path[-1] == base_address
                amount_out = int(qty / TOKEN_TO_DECIMAL_MULTIPLIER[base_currency])
                amount_in_maximum = int(
                    qty * price * (1 + slippage) / TOKEN_TO_DECIMAL_MULTIPLIER[quote_currency])
                func = self.router.functions.swapTokensForExactTokens

                if args is None:
                    args = (amount_out, amount_in_maximum, path, wallet_address, deadline)
            case 'SELL':
                # reform swap path, short base and long quote
                if path[0] != base_address:
                    path = path[::-1]  # swap in base and swap out quote
                assert path[0] == base_address
                assert path[-1] == quote_address
                amount_in = int(qty / TOKEN_TO_DECIMAL_MULTIPLIER[base_currency])
                amount_out_minimum = int(
                    qty * price * (1 - slippage) / TOKEN_TO_DECIMAL_MULTIPLIER[quote_currency])
                # SupportingFeeOnTransferTokens is for CCX case
                func = self.router.functions.swapExactTokensForTokensSupportingFeeOnTransferTokens

                if args is None:
                    args = (amount_in, amount_out_minimum, path, wallet_address, deadline)
            case _:
                raise NotImplementedError()

        tx_param = func(*args).build_transaction(
            {
                "from": wallet_address,
                "value": ETH_USED,
            }
        ) | tx_param
        logging.info('before sign: %s' % tx_param)
        return sign_and_wait(self._web3, [tx_param], symbol, self._blockchain_name)

    def approve_submit_order(
        self,
        wallet_address,
        symbol,
    ) -> bool:
        return self.approve_SwapRouter(wallet_address, symbol)

    def approve_SwapRouter(
        self,
        wallet_address: str,
        symbol: str,
        amount=MAX_UINT_256,
    ) -> bool:
        p = UniswapProduct.FromStr(symbol)
        currencies = []
        if p.base.currency == 'ETH':
            currencies.append('WETH')
        else:
            currencies.append(p.base.currency)
        if p.quote.currency == 'ETH':
            currencies.append('WETH')
        else:
            currencies.append(p.quote.currency)
        contracts = []
        for currency in currencies:
            if currency == 'USDT':
                contract_addr = TOKEN_TO_ADDRESS['ethereum']['USDT']
                contract = _load_contract(self._web3, "usdt", Web3.to_checksum_address(contract_addr))
            else:
                contract_addr = TOKEN_TO_ADDRESS['ethereum'][currency]
                contract = _load_contract(self._web3, "erc20", Web3.to_checksum_address(contract_addr))
            contracts.append(contract)

        for contract in contracts:
            tx_param = contract.functions.approve(
                deployment_address["UniswapV2Router02"], amount).build_transaction(
                {
                    "from": wallet_address,
                    "type": "0x2",
                    "maxFeePerGas": self._web3.eth.max_priority_fee + (
                        2 * self._web3.eth.get_block("latest")["baseFeePerGas"]),
                    "maxPriorityFeePerGas": self._web3.eth.max_priority_fee,
                    "chainId": BLOCKCHAIN_NAME_TO_CHAINID[self._blockchain_name],
                }
            )
            tx_param["gas"] = int(self._web3.eth.estimate_gas(tx_param) * 1.2)
            logging.info('before sign: %s' % tx_param)
            ret = sign_and_wait(self._web3, [tx_param], None, self._blockchain_name)
            if not ret[0]:
                return False
        return True
